package com.abc.config.memory;

import cn.hutool.core.util.ObjectUtil;
import com.abc.enums.IdType;
import com.abc.help.SessionHelper;
import com.abc.model.ChatMessageModel;
import com.abc.service.IdGeneratorService;
import com.abc.util.MessageChainUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.Sort;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@Component
public class MongodbChatMemoryRepository implements ChatMemoryRepository {

    @Autowired
    private MongoTemplate mongoTemplate;

    @Autowired
    private IdGeneratorService idGeneratorService;
    
    @Autowired
    private MessageChainUtil messageChainUtil;


    @Override
    public List<String> findConversationIds() {
        Query query = Query.query(Criteria.where("session_id").exists(true).and("username").is(SessionHelper.getUser()));
        return mongoTemplate.find(query, ChatMessageModel.class).stream().map(ChatMessageModel::getSessionId).collect(Collectors.toList());
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        Query query = Query.query(Criteria.where("session_id").is(conversationId).and("username").is(SessionHelper.getUser()))
                .with(Sort.by(Sort.Direction.ASC, "create_time"));
        return mongoTemplate.find(query, ChatMessageModel.class).stream()
                .map(ChatMessageModel::toMessage)
                .collect(Collectors.toList());
    }
    
    /**
     * 根据对话ID查找最新的指定数量的消息
     *
     * @param conversationId 对话ID
     * @param limit 限制数量
     * @return 消息列表
     */
    public List<Message> findLatestByConversationId(String conversationId, int limit) {
        Query query = Query.query(Criteria.where("session_id").is(conversationId).and("username").is(SessionHelper.getUser()))
                .with(Sort.by(Sort.Direction.DESC, "create_time"))
                .limit(limit);
        
        List<ChatMessageModel> models = mongoTemplate.find(query, ChatMessageModel.class);
        // 反转列表，使其按时间正序排列
        Collections.reverse(models);
        
        return models.stream()
                .map(ChatMessageModel::toMessage)
                .collect(Collectors.toList());
    }

    /**
     * 根据对话ID查找最新的指定数量的消息模型
     *
     * @param conversationId 对话ID
     * @param limit 限制数量
     * @return ChatMessageModel列表
     */
    public List<ChatMessageModel> findLatestMessageModelsByConversationId(String conversationId, int limit) {
        Query query = Query.query(Criteria.where("session_id").is(conversationId).and("username").is(SessionHelper.getUser()))
                .with(Sort.by(Sort.Direction.DESC, "create_time"))
                .limit(limit);
        
        List<ChatMessageModel> models = mongoTemplate.find(query, ChatMessageModel.class);
        // 反转列表，使其按时间正序排列
        Collections.reverse(models);
        
        return models;
    }


    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        List<ChatMessageModel> savedModels = saveAllAndReturn(conversationId, messages);
        log.debug("保存{}条消息到会话: {}", savedModels.size(), conversationId);
    }

    /**
     * 保存消息并返回生成的ChatMessageModel列表
     *
     * @param conversationId 对话ID
     * @param messages 消息列表
     * @return 保存的ChatMessageModel列表
     */
    public List<ChatMessageModel> saveAllAndReturn(String conversationId, List<Message> messages) {
        if (ObjectUtil.isEmpty(messages)) {
            return new ArrayList<>();
        }
        
        String username = SessionHelper.getUser();
        
        // 获取会话中最后一条消息的ID，作为新消息的父消息
        String lastMessageId = messageChainUtil.getLastMessageId(conversationId, username);
        
        // 批量生成消息ID，确保连续性
        List<String> messageIds = idGeneratorService.generateIds(IdType.MESSAGE, messages.size());
        
        // 创建ChatMessageModel列表
        List<ChatMessageModel> chatMessageModels = new ArrayList<>();
        String parentMessageId = lastMessageId; // 第一条新消息的父消息是最后一条已存在的消息
        
        for (int i = 0; i < messages.size(); i++) {
            Message message = messages.get(i);
            String messageId = messageIds.get(i);
            
            // 构建当前消息的链路
            List<String> messageChain = messageChainUtil.buildMessageChainFromParent(
                    conversationId, username, parentMessageId, messageId);
            
            ChatMessageModel model = new ChatMessageModel(
                    conversationId,
                    message.getMessageType().getValue(),
                    username,
                    message.getText(),
                    messageId,
                    message.getMetadata(),
                    parentMessageId, // 设置父消息ID
                    messageChain     // 设置消息链路
            );
            chatMessageModels.add(model);
            
            // 下一条消息的父消息是当前消息
            parentMessageId = messageId;
        }
        
        // 保存新消息
        mongoTemplate.insertAll(chatMessageModels);
        
        return chatMessageModels;
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        log.warn("持久化存储不支持删除会话记录");
    }
}
